#include "Communication.hpp"
#include <iterator>
#include <sstream>
#include <iostream>
#include <cstring>
#include <cstdlib>
#include <cstdio>

extern "C"{
#include <arpa/inet.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netdb.h>
#include <errno.h>
}

static ostream &cerrerr()
{
	return cerr << "ERROR: ";
}

static void test(int cond, const char* str)
{
	if(cond < 0){
		cerrerr();
		perror(str);
		exit(-1);
	}
}

int whoami(list<string> name)
{
	size_t length = 0;
	list<string>::iterator it;
	for(it = name.begin(); it != name.end(); it++)
		length = length > it->size() ? length : it->size();
	length++;
	char *me = new char[length];
	int res = -1;
	if(gethostname(me, length)){
		res = -1;
		cerr << "whoami: Cannot get hostname" << endl;
	}else{
		cerr << "whoami: Hostname is " << me << endl;
		it = name.begin();
		for(unsigned i = 0; i < name.size(); i++, it++)
			if(*it == me){
				res = i;
				cerr << "whoami: Host id is " << i << endl;
				break;
			}
	}
	delete [] me;
	return res;
}

void LogicalClock::operator+=(const LogicalClock &that)
{
	if(size() != that.size()){
		cerrerr() << "LogicalClock::operator+=: Size mismatch" << endl;
		exit(-1);
	}
	for(unsigned i = 0; i < size(); i++)
		if(that[i] > (*this)[i])
			(*this)[i] = that[i];
}

LogicalClock::LogicalClock(string str)
{
	clear();
	istringstream in(str);
	char c;
	int n;
	enum{s_begin, s_prenum, s_postnum, s_end} stat = s_begin;
	while(stat != s_end)
		switch(stat){
		case s_begin:
			in >> c;
			n = in.peek();
			if(in.fail() || c != '<'){
				cerrerr() << "LogicalClock::operator string: Unexpected string head: " << str << endl;
				exit(-1);
			}
			stat = n == '>' ? s_postnum : s_prenum;
			break;
		case s_prenum:
			in >> n;
			if(in.fail()){
				cerrerr() << "LogicalClock::operator string: Unexpected string content" << str << endl;
				exit(-1);
			}
			push_back(n);
			stat = s_postnum;
			break;
		case s_postnum:
			in >> c;
			if(in.fail()){
				cerrerr() << "LogicalClock::operator string: Unexpected string delimiter: " << str << endl;
				exit(-1);
			}
			switch(c){
			case ',':
				stat = s_prenum;
				break;
			case '>':
				in >> c;
				if(!in.eof()){
					cerrerr() << "LogicalClock::operator string: Unexpected string ending: " << str << endl;
					exit(-1);
				}
				stat = s_end;
				break;
			default: cerrerr() << "LogicalClock::operator string: Unexpected string: " << str << endl;
			}
			break;
		default:
			cerrerr() << "LogicalClock::operator string: Unknown status" << stat << endl;
		}
}

LogicalClock::operator string(void) const
{
	ostringstream out;
	out << '<';
	for(unsigned i = 0; i < size(); i++)
		out << (*this)[i] << ',';
	if(size() > 0)
		out.seekp(-1, ios_base::cur);
	out << '>';
	return out.str();
}

Message::operator string(void) const
{
	ostringstream out;
	if(deliv > 0)
		out << deliv << sep[0];
	out << (string)ts << sep[1] << data;
	return out.str();
}

Message::Message(const string &str)
{
	*this = Message();
	size_t pos = str.find_first_of(sep), pos2;
	if(pos == str.npos){
		cerrerr() << "Message::Message: Unexpected string: " << str << endl;
		exit(-1);
	}else if(str[pos] == sep[0]){
		istringstream in(str.substr(0, pos));
		in >> deliv;
		if(in.fail()){
			cerrerr() << "Message::Message: Unexpected string: " << str << endl;
			exit(-1);
		}
		char c;
		in >> c;
		if(!in.eof()){
			cerrerr() << "Message::Message: Unexpected string: " << str << endl;
			exit(-1);
		}
		pos2 = str.find_first_of(sep[1], ++pos);
		if(pos2 == str.npos){
			cerrerr() << "Message::Message: Unexpected string: " << str << endl;
			exit(-1);
		}
	}else if(str[pos] == sep[1]){
		deliv = 0;
		pos2 = pos;
		pos = 0;
	}else{
		cerrerr() << "Message::Message: Unexpected string: " << str << endl;
		exit(-1);
	}
	ts = LogicalClock(str.substr(pos,  pos2 - pos));
	data = str.substr(pos2 + 1);
}

static char* getaddr(const struct sockaddr *sa)
{
	static char str[INET6_ADDRSTRLEN + 6];
	void *p;
	unsigned short port;
	switch(sa->sa_family){
	case AF_INET:
		p = &(((struct sockaddr_in *)sa)->sin_addr);
		port = ((struct sockaddr_in *)sa)->sin_port;
		break;
	case AF_INET6:
		p = &(((struct sockaddr_in6 *)sa)->sin6_addr);
		port = ((struct sockaddr_in6 *)sa)->sin6_port;
		break;
	default:
		cerrerr() << "getaddr: Unknown sa_family " << sa->sa_family << endl;
		exit(-1);
	}
	test(inet_ntop(sa->sa_family, p, str, INET6_ADDRSTRLEN) == NULL ? -1 : 0, "getaddr");
	port = ntohs(port);
	sprintf(str + strlen(str), ":%d", port);
	return str;
}

static char* getaddr(const struct addrinfo *ai)
{
	return getaddr(ai->ai_addr);
}

Communication::Communication(list<Host> allhost, bool issrv, int me)
{
	// to count message received from socket in every time slot
	count = 0;
	last = 0;
	// self id and vc
	num = allhost.size();
	if(me < 0){
		list<string> name;
		for(list<Host>::iterator it = allhost.begin(); it != allhost.end(); it++)
			name.push_back(it->name);
		self = whoami(name);
	}else
		self = me;
	if(self < 0){
		cerrerr() << "Communication::Communication::whoami: Not in name list" << endl;
		exit(-1);
	}
	vc = LogicalClock(num, self);
	// hints
	struct addrinfo hints;
	memset(&hints, 0, sizeof(hints));
	hints.ai_family = AF_INET;
	hints.ai_socktype = SOCK_STREAM;
	hints.ai_flags = AI_PASSIVE;
	// initialize sockets and addrinfo
	sock.resize(num, 0);
	struct addrinfo **res = new struct addrinfo* [num];
	int err;
	char p[6];
	list<Host>::iterator it = allhost.begin();
	for(int i = 0; i < num; i++, it++){
		if(it->port < 1024 || it->port > 65535){
			cerrerr() << "Communication::Communication: Invalid port " << it->port << endl;
			exit(-1);
		}
		sprintf(p, "%d", it->port);
		if((err = getaddrinfo(it->name.c_str(), p, &hints, &res[i]))){
			cerrerr() << "Communication::Communication::getaddrinfo: " << gai_strerror(err) << endl;
			exit(-1);
		}else
			cerr << "Communication::Communication: Address of node " << i << ": " << getaddr(res[i]) << endl;
	}
	// get connections
	if(issrv){ // if it is the server
		cerr << "Communication::Communication: I'm the server" << endl;
		server(res);
	}
	else{ // if it is a peer process
		cerr << "Communication::Communication: I'm peer node " << self << endl;
		peer(res);
	}
	cerr << "Communication::Communication: All connections ready" << endl;
	for(int i = 0; i < num; i++)
		freeaddrinfo(res[i]);
	delete [] res;
}

void Communication::server(struct addrinfo **ai)
{
	srv = 0;
	// establish connection to every peer
	recvconn(ai, num - 1);
}

void Communication::peer(struct addrinfo **ai)
{
	// establish connection to server
	srv = sendconn(ai[self], ai[0]);
	// establish connection to peers after me
	for(int i = self + 1; i < num; i++)
		sock[i] = sendconn(ai[self], ai[i]);
	// receive connection from peers before me
	if(self > 0)
		recvconn(ai, self - 1);
	else
		cerr << "Communicaion::peer: No connection to receive" << endl;
}

int Communication::sendconn(const struct addrinfo *me, struct addrinfo *target)
{
	int s = socket(me->ai_family, me->ai_socktype, 0);
	test(s, "Communication::sendconn::socket");
	int err = 0;
	while((err = connect(s, target->ai_addr, target->ai_addrlen)))
		if(errno == ECONNREFUSED)
			sleep(2);
		else
			test(err, "Communication::sendconn::connect");
	test(::send(s, &self, sizeof(self), 0), "Communication::sendconn::send");
	cerr << "Communication::sendconn: Connected to " << getaddr(target) << endl;
	return s;
}

void Communication::recvconn(struct addrinfo **ai, int to)
{
	// create listening socket
	int l = socket(ai[self]->ai_family, ai[self]->ai_socktype, 0);
	test(l, "Communication::recvconn::socket");
	int err;
	while((err = bind(l, ai[self]->ai_addr, ai[self]->ai_addrlen)))
		if(errno == EADDRINUSE)
			sleep(2);
		else
			test(err, "Communication::recvconn::bind");
	test(listen(l, to + 1), "Communication::recvconn::listen");
	cerr << "Communication::recvconn: Listening on " << getaddr(ai[self]) << endl;
	// accept all incoming connections
	struct sockaddr inaddr;
	//inaddr.ss_family = ai[self]->ai_family;
	socklen_t insize = sizeof(inaddr);
	int s;
	int from;
	bool match;
	for(int i = 0; i < to + 1; i++){
		s = accept(l, &inaddr, &insize);
		test(s, "Communication::recvconn::accept");
		test(::recv(s, &from, sizeof(from), MSG_WAITALL), "Communication::recvconn::recv");
		if(from < 0 || from >= num){
			cerrerr() << "Communication::recvconn: Invalid incoming node number " << from << endl;
			exit(-1);
		}
		if(srv > 0 && from == self){
			cerrerr() << "Communication::recvconn: Connection from self" << endl;
			exit(-1);
		}
		if(sock[from] > 0){ // if the socket is taken
			cerrerr() << "Communication::recvconn: Duplicate connection from node " << from << endl;
			exit(-1);
		}
		switch(ai[from]->ai_family){
		case AF_INET:
			match = ((struct sockaddr_in *)&inaddr)->sin_addr.s_addr == ((struct sockaddr_in *)ai[from]->ai_addr)->sin_addr.s_addr;
			break;
		case AF_INET6:
			match = true;
			for(int k = 0; k < 16; k++)
				if(((struct sockaddr_in6 *)&inaddr)->sin6_addr.s6_addr != ((struct sockaddr_in6 *)ai[from]->ai_addr)->sin6_addr.s6_addr)
					match = false;
			break;
		default:
			cerrerr() << "Communication::recvconn: Unknown ai_family " << ai[from]->ai_family << endl;
		}
		if(!match){
			cerrerr() << "Communication::recvconn: Unknown incomming connection from node " << from << " " << getaddr(&inaddr) << endl;
			exit(-1);
		}
		cerr << "Communication::recvconn: Received connection from node " << from << " " << getaddr(&inaddr) << endl;
		sock[from] = s;
	}
	cerr << "Communication::recvconn: All connection received." << endl;
	close(l); // close listening socket
}

void Communication::send(Message &msg)
{
	const char *prefix = "Communication::send::send";
	// timestamping
	if(srv > 0 && msg.node > -2){ // if the message is not from/to a server
		vc++;
		msg.ts = vc;
		//cerr << "Communication::send: Update clock to " << (string)vc << endl;
	}else
		msg.ts = LogicalClock();
	// serialization, insert sperator and length
	string str = (string) msg;
	if(str.size() > 0xFFFF){
		cerrerr() << "Communication::send: Message too long " << str.size();
		exit(-1);
	}
	char size[2] = {str.size() % 256, str.size() / 256};
	str.insert(0, msg.sep);
	str.insert(2, 1, size[0]);
	str.insert(3, 1, size[1]);
	// send message
	switch(msg.node){
	case -2: // to server
		//cerr << "Communication::send: Message to server" << endl;
		if(srv > 0)
			test(::send(srv, str.c_str(), str.size(), 0), prefix);
		else{
			cerrerr() << "Communication::send: Server cannot send to server" << endl;
			exit(-1);
		}
		break;
	case -1: // broadcast
		//cerr << "Communication::send: Message to everyone" << endl;
		for(int i = 0; i < num; i++)
			if(sock[i] > 0)
				test(::send(sock[i], str.c_str(), str.size(), 0), prefix);
		break;
	default: // to peer
		//cerr << "Communication::send: Message to node " << msg.node << endl;
		if(msg.node < 0 || msg.node > num - 1 || msg.node == self){
			cerrerr() << "Communication::send: No socket to node " << msg.node << endl;
			exit(-1);
		}
		test(::send(sock[msg.node], str.c_str(), str.size(), 0), prefix);
	}
	//cerr << "Communication::send: Message sent: " << str.substr(4) << endl;
}

int Communication::recv(Message &msg, unsigned now)
{
	//cerr << "Communication::recv: Now is " << now << endl;
	if(now != last){
		count = 0;
		last = now;
	}
	// check if any socket has new data
	fd_set set;
	struct timeval tv;
	FD_ZERO(&set);
	vector<int> ss = sock;
	if(srv > 0)
		ss[self] = srv;
	int max = 0;
	for(int i = 0; i < num; i++){
		FD_SET(ss[i], &set);
		max = max > ss[i] ? max : ss[i];
	}
	tv.tv_sec = 0;
	tv.tv_usec = 0;
	int ready;
	test(ready = select(max + 1, &set, NULL, NULL, &tv), "Communication::recv::select");
	//cerr << "Communication::recv: Ready socket number " << ready << endl;
	// check messages in channels
	Message m;
	if(ready > 0){
		int size, rs, src;
		bool next;
		unsigned char *buffer = (unsigned char*) malloc(4);
		for(int i = 0; i < num; i++)
			if(FD_ISSET(ss[i], &set)){ // check if the socket is ready to read
				src = srv > 0 && i == self ? -2 : i; // source node id
				//cerr << "Communication::recv: Get message from node " << src << endl;
				next = true;
				while(next){
					// peek and process the message header
					test(rs = ::recv(ss[i], buffer, 4, MSG_PEEK), "Communication::recv::recv");
					if(rs < 4){ // if header incomplete, skip the socket
						//cerr << "Communication::recv: Header incomplete " << rs << endl;
						break;
					}
					if(buffer[0] != m.sep[0] || buffer[1] != m.sep[1]){ // if header unexpected
						cerrerr() << "Communication::recv: Message boundary mismatch" << endl;
						exit(-1);
					}
					size = buffer[2] + buffer[3] * 256 + 4; // size of the whole header+message
					//cerr << "Communication::recv: Message size " << size - 4 << endl;
					buffer = (unsigned char*) realloc(buffer, size + 1); // one more byte to test if next message ready
					// peek and test the header+message
					test(rs = ::recv(ss[i], buffer, size + 1, MSG_PEEK), "Communication::recv::recv");
					if(rs < size){ // if message incomplete, skip the socket
						//cerr << "Communication::recv: Message incomplete " << rs - 4 << endl;
						break;
					}else if(rs == size) // if next message is not received
						next = false;
					// retrieve the header+message
					test(rs = ::recv(ss[i], buffer, size, 0), "Communication::recv::recv");
					if(rs < size){ // if message incomplete, something is wrong
						cerrerr() << "Communication::recv: Message too short" << endl;
						exit(-1);
					}
					m = Message(string((const char*)buffer + 4, size - 4)); // convert string to message
					m.node = src;
					//cerr << "Communication::recv: Message received: " << string((char*)buffer + 4, size - 4) << endl;
					msgq.push(m); // enqueue
					if(src >= 0)
						count++;
				}
			}
		free(buffer);
	}
	// check messages in queue
	if(msgq.size() == 0)
		return msgq.size();
	m = msgq.top(); // peek the top message
	//cerr << "Communication::recv: Next delivery at " << m.deliv << endl;
	if(m.deliv > now)
		return msgq.size();
	// handle the message to deliver
	msgq.pop();
	msg = m;
	if(msg.ts.size() > 0){ // update clock
		if(msg.node < 0 || srv == 0){
			cerrerr() << "Communication::recv: Message to/from server shouldn't be timestamped" << endl;
			exit(-1);
		}
		vc++;
		vc += msg.ts;
		//cerr << "Communication::recv: Update clock to " << (string) vc << endl;
	}
	//cerr << "Communication::recv: Message delivered: " << (string) msg << endl;
	return - 1 - msgq.size();
}
